import unittest
from AStock.ASQuery import (
    ASQuery_stock_info,
    ASQuery_research_report,
    ASQuery_stock_day,
    ASQuery_valuation,
    ASQuery_trade_cal,
    ASQuery_financial,
    ASQuery_nearest_trade_date,
    ASQuery_stock_adj,
    ASQuery_stock_xdxr
)


class TestQuery(unittest.TestCase):

    def testQueryStockInfo(self):
        df = ASQuery_stock_info('300001')
        print(df.head())
        self.assertEqual(df.shape[0], 1)
        self.assertEqual(df.loc[0, 'code'], '300001')

        df = ASQuery_stock_info()
        print(df.shape)
        print(df.head(10))

        df = ASQuery_stock_info(['300001', '000001'])
        print(df.head())
        self.assertEqual(df.shape[0], 2)
        self.assertEqual(df.loc[0, 'code'], '000001')
        self.assertEqual(df.loc[1, 'code'], '300001')

    def testQueryResearchReport(self):
        df = ASQuery_research_report(
            '600570',
            '20200317',
            '20200917'
        )
        print(df.head(10))
        print(df.shape)

        df = ASQuery_research_report(
            '002511',
            '20200617',
            '20200917',
            fields=['stockCode', 'stockName', 'predictThisYearEps', 'publishDate'],
            limit=10
        )
        print(df.head(10))
        print(df.shape)
        self.assertEqual(df.shape, (10, 4))

    def testQueryStockDay(self):
        df = ASQuery_stock_day('600570', '20200910', '20200917', fields=['code', 'close', 'date'])
        print(df.head(10))

        df = ASQuery_stock_day('600585', '20200331', '20200331', type='qfq')
        print(df.head(10))

        df = ASQuery_stock_day('600585', '20200331', '20200331', fields=['code', 'close', 'total_mv', 'circ_mv'])
        print(df.head(10))

    def testQueryValuation(self):
        df = ASQuery_valuation(code='600570')
        print(df.head(10))

        df = ASQuery_valuation(fields=['code', 'pe30'])
        print(df.head(10))
        print(df.shape)

    def testQueryTradeCal(self):
        self.assertTrue(ASQuery_trade_cal('20200924'))
        self.assertFalse(ASQuery_trade_cal('20200926'))

    def testQueryNearestTradeDate(self):
        self.assertTrue('20200924' == ASQuery_nearest_trade_date('20200924'))
        self.assertTrue('20200925' == ASQuery_nearest_trade_date('20200927'))
        self.assertTrue('20200928' == ASQuery_nearest_trade_date('20200927', future=True))
        self.assertIsNone(ASQuery_nearest_trade_date('20200927', max_days=1))

    def testQueryFinancial(self):
        df = ASQuery_financial(code='300003', report_date='20200630', convert_column_name=True)
        self.assertFalse(df.empty)
        print(df.head(10))

        df = ASQuery_financial(code='300003', report_date='20200630', fields=['code', 'report_date', '228', '281', 'test'])
        self.assertFalse(df.empty)
        print(df.head(10))

    def testQueryStockAdj(self):
        df = ASQuery_stock_adj('300003', '20200630', '20200710')
        self.assertEqual(df.shape[0], 9)
        print(df)
        df = ASQuery_stock_adj(['000001', '300003'], '20200630', '20200710')
        self.assertEqual(df.shape[0], 18)
        print(df)

    def testQueryStockXdxr(self):
        df = ASQuery_stock_xdxr('300003')
        print(df.shape)
        print(df)
        df = ASQuery_stock_xdxr(['000001', '300003'], fields=['name', 'code'])
        print(df.shape)
        print(df)